import os
import sys
import time
import paramiko
from paramiko import SFTPServerInterface

from .connection import SSHConnection
from .const import PERMS_ACTION_NAME_DOWNLOAD_FILE, PERMS_ACTION_NAME_UPLOAD_FILE, PERMS_ACTION_NAME_CONNECT
from apps.extensions import db
from apps.assets.models import Asset, SystemUser
from apps.terminal.models import SFTPLog


CURRENT_DIR = os.path.dirname(__file__)


class SFTPServer(SFTPServerInterface):
    ROOT = '/home'

    def __init__(self, server, **kwargs):
        super(SFTPServer, self).__init__(server, **kwargs)
        self.server = server
        self._sftp = {}
        self.hosts = self.get_permed_hosts()
        self.is_finished = False

    @staticmethod
    def get_host_and_user_object(host, user):
        # todo 需要增加缓存
        asset = db.session.query(Asset).filter(Asset.hostname == host).first()
        system_user = db.session.query(SystemUser).filter(SystemUser.username == user).first()
        return asset, system_user

    def get_permed_hosts(self):
        """
        :return:
            {
                'testServer02': ['sa', 'dev'],
                'testServer03': ['sa', 'dev']
            }
        """

        hosts = {}
        current_user = self.server.connection.user
        results = current_user.get_related_assets()
        for _, data in results.items():
            asset = data.get('asset')
            system_users = data.get('system_users')
            for system_user in system_users:
                hostname = asset.hostname
                if hostname not in hosts.keys():
                    hosts[hostname] = []
                hosts[hostname].append(system_user.username)
        return hosts

    def get_host_sftp(self, host, user):
        asset, system_user = self.get_host_and_user_object(host, user)
        if not asset or not system_user:
            raise PermissionError("No asset or system user explicit")

        cache_key = '{}@{}'.format(user, host)
        if cache_key not in self._sftp:
            conn = SSHConnection.new_connection(self.server.connection.user,
                                                asset, system_user)
            __sftp = conn.get_sftp()
            if __sftp:
                sftp = {
                    'client': __sftp, 'connection': conn,
                    'home': __sftp.normalize('')
                }
                self._sftp[cache_key] = sftp
                return sftp
            else:
                raise OSError("Can not connect asset sftp service: {}".format(conn.error))
        else:
            return self._sftp[cache_key]

    def session_ended(self):
        self.is_finished = True
        super(SFTPServer, self).session_ended()
        for _, v in self._sftp.items():
            sftp = v['client']
            conn = v.get('connection')
            sftp.close()
            conn.close()
        self._sftp = {}

    @staticmethod
    def parse_path(path):
        """
            /hostname/system_user/path
            /<host>/<user>/<dpath>
        :param path:
        :return:
        """
        data = path.lstrip('/').split('/')
        request = {"host": "", "user": "", "dpath": ""}

        if len(data) == 1 and not data[0]:
            return request

        host, path = data[0], data[1:]
        request['host'] = host
        request['user'], path = (path[0], path[1:]) if path else ('', path)
        request['dpath'] = '/'.join(path)
        return request

    def create_ftp_log(self, path, operate, is_success=True, filename=None):
        request = self.parse_path(path)
        host, user = request['host'], request['user']
        c, rpath = self.get_sftp_client_rpath(request)
        data = {
            "user": "{} ({})".format(self.server.connection.user.name,
                                     self.server.connection.user.username),
            "asset": host,
            "system_user": user,
            "remote_addr": self.server.connection.addr[0],
            "operate": operate,
            "filename": filename or rpath,
            "is_success": is_success,
            "create_time": time.time()
        }
        SFTPLog.create(**data)

    def get_sftp_client_rpath(self, request):
        self.check_action(request, PERMS_ACTION_NAME_CONNECT)
        sftp = self.get_host_sftp(request['host'], request['user'])
        rpath = os.path.join(self.ROOT, request['user'], request['dpath'].lstrip('/'))
        return sftp['client'], rpath

    def validate_permission(self, host, user, action):
        asset, system_user = self.get_host_and_user_object(host, user)
        current_user = self.server.connection.user
        return current_user.has_asset_action_perm(asset, system_user, action)

    def check_action(self, request, action):
        host, user = request['host'], request['user']
        if not self.validate_permission(host, user, action):
            raise PermissionError("Permission deny")

    @staticmethod
    def stat_fake_dir():
        s = os.stat(CURRENT_DIR)
        attr = paramiko.SFTPAttributes.from_stat(s)
        return attr

    def list_folder(self, path):
        output = []
        request = self.parse_path(path)
        print("List folder: {} => {}".format(path, request))
        if not request['host']:  # It's root(根目录), 显示用户的所有资产
            for hostname in self.hosts.keys():
                attr = self.stat_fake_dir()
                attr.filename = hostname
                output.append(attr)
        elif not request['user']:  # It's asset(资产目录), 显示资产的所有用户
            for user in self.hosts.get(request['host']):
                attr = self.stat_fake_dir()
                attr.filename = user
                output.append(attr)
        else:
            client, rpath = self.get_sftp_client_rpath(request)
            output = client.listdir_attr(rpath)
            output = [attr for attr in output if not attr.filename.startswith('.')]
        return output

    def stat(self, path):
        request = self.parse_path(path)
        host, user, dpath = request['host'], request['user'], request['dpath']

        print("Stat path: {} => {}".format(path, request))
        if not host or not user:
            stat = self.stat_fake_dir()
            stat.filename = host or user or '/'
            return stat

        client, rpath = self.get_sftp_client_rpath(request)
        stat = client.stat(rpath)
        return stat

    def lstat(self, path):
        return self.stat(path)

    def open(self, path, flags, attr):
        binary_flag = getattr(os, 'O_BINARY', 0)
        flags |= binary_flag
        success = False

        if flags & os.O_WRONLY:
            if flags & os.O_APPEND:
                mode = 'ab'
            else:
                mode = 'wb'
        elif flags & os.O_RDWR:
            if flags & os.O_APPEND:
                mode = 'a+b'
            else:
                mode = 'r+b'
        else:
            mode = 'rb'

        if 'r' in mode:
            operate = "Download"
            action = PERMS_ACTION_NAME_DOWNLOAD_FILE
        elif 'a' in mode:
            operate = "Append"
            action = PERMS_ACTION_NAME_UPLOAD_FILE
        else:
            operate = "Upload"
            action = PERMS_ACTION_NAME_UPLOAD_FILE

        try:
            request = self.parse_path(path)
            self.check_action(request, action)

            client, rpath = self.get_sftp_client_rpath(request)
            f = client.open(rpath, mode, bufsize=4096)
            f.prefetch()
            obj = paramiko.SFTPHandle(flags)
            obj.filename = rpath
            obj.readfile = f
            obj.writefile = f
            result = obj
            success = True
            return result
        finally:
            self.create_ftp_log(path, operate, success)

    def remove(self, path):
        request = self.parse_path(path)
        self.check_action(request, action=PERMS_ACTION_NAME_UPLOAD_FILE)

        client, rpath = self.get_sftp_client_rpath(request)
        success = False

        try:
            client.remove(rpath)
            success = True
            return paramiko.SFTP_OK
        finally:
            self.create_ftp_log(path, "Delete", success)

    def rename(self, src, dest):
        src_request = self.parse_path(src)
        self.check_action(src_request, action=PERMS_ACTION_NAME_UPLOAD_FILE)

        dest_request = self.parse_path(dest)
        client, rsrc = self.get_sftp_client_rpath(src_request)
        client2, rdest = self.get_sftp_client_rpath(dest_request)
        success = False
        filename = "{}=>{}".format(rsrc, rdest)
        try:
            if client == client2:
                client.rename(rsrc, rdest)
                success = True
                return paramiko.SFTP_OK
            else:
                raise FileNotFoundError("Can't rename {} => {}".format(src, dest))
        finally:
            self.create_ftp_log(src, "Rename", success, filename=filename)

    def mkdir(self, path, attr=0o755):
        request = self.parse_path(path)
        self.check_action(request, action=PERMS_ACTION_NAME_UPLOAD_FILE)

        client, rpath = self.get_sftp_client_rpath(request)
        success = False

        try:
            if rpath == '/':
                raise PermissionError("Create '/', are you sure?")
            client.mkdir(rpath)
            success = True
            return paramiko.SFTP_OK
        finally:
            self.create_ftp_log(path, "Mkdir", success)

    def rmdir(self, path):
        request = self.parse_path(path)
        self.check_action(request, action=PERMS_ACTION_NAME_UPLOAD_FILE)

        client, rpath = self.get_sftp_client_rpath(request)
        success = False

        try:
            if rpath == '/':
                raise PermissionError("Delete '/', are you sure?")
            client.rmdir(rpath)
            success = True
            return paramiko.SFTP_OK
        finally:
            self.create_ftp_log(path, "Rmdir", success)

    def chattr(self, path, attr):
        request = self.parse_path(path)
        client, rpath = self.get_sftp_client_rpath(request)
        success = False

        try:
            if sys.platform != 'win32':
                # mode operations are meaningless on win32
                if attr._flags & attr.FLAG_PERMISSIONS:
                    client.chmod(rpath, attr.st_mode)
                if attr._flags & attr.FLAG_UIDGID:
                    client.chown(rpath, attr.st_uid, attr.st_gid)
            if attr._flags & attr.FLAG_AMTIME:
                client.utime(rpath, (attr.st_atime, attr.st_mtime))
            if attr._flags & attr.FLAG_SIZE:
                with open(rpath, 'w+') as f:
                    f.truncate(attr.st_size)

            success = True
            return paramiko.SFTP_OK
        finally:
            self.create_ftp_log(path, "Chattr", success)

    def symlink(self, target_path, original_path):
        target_path_request = self.parse_path(target_path)
        original_path_request = self.parse_path(original_path)

        client, target_path = self.get_sftp_client_rpath(target_path_request)
        client2, path = self.get_sftp_client_rpath(original_path_request)

        success = False
        filename = "{}=>{}".format(target_path, path)

        try:
            if client == client2:
                client.symlink(target_path, path)
                success = True
                return paramiko.SFTP_OK
            else:
                raise FileNotFoundError("Can't symlink {} => {}".format(target_path, path))
        finally:
            self.create_ftp_log(target_path, "Symlink ", success, filename=filename)

    def readlink(self, path):
        request = self.parse_path(path)
        client, rpath = self.get_sftp_client_rpath(request)
        symlink = client.readlink(rpath)
        return symlink


class FakeServer:
    pass


class FakeTransport:
    _trans = None

    @staticmethod
    def getpeername():
        return '127.0.0.1', 6666

    @staticmethod
    def get_username():
        return 'fake'


class FakeChannel:
    _chan = None

    def get_transport(self):
        return FakeTransport()

    @classmethod
    def new(cls):
        if not cls._chan:
            cls._chan = cls()
        return cls._chan


class InternalSFTPClient(SFTPServer):
    def __init__(self, connection):
        fake_server = FakeServer()
        fake_server.connection = connection
        super(InternalSFTPClient, self).__init__(fake_server)

    def listdir_attr(self, path):
        return self.list_folder(path)

    def open(self, path, mode, **kwargs):
        request = self.parse_path(path)
        client, rpath = self.get_sftp_client_rpath(request)
        if 'r' in mode:
            operate = "Download"
            action = PERMS_ACTION_NAME_DOWNLOAD_FILE
        else:
            operate = "Upload"
            action = PERMS_ACTION_NAME_UPLOAD_FILE
        success = False
        try:
            self.check_action(request, action=action)
            f = client.open(rpath, mode, bufsize=4096)
            success = True
            return f
        finally:
            self.create_ftp_log(path, operate, success)

    def stat(self, path):
        attr = super(InternalSFTPClient, self).stat(path)
        return attr

    def lstat(self, path):
        attr = super(InternalSFTPClient, self).lstat(path)
        return attr

    def rmdir(self, path):
        return super(InternalSFTPClient, self).rmdir(path)

    def get_channel(self):
        return FakeChannel.new()

    def unlink(self, path):
        return self.remove(path)

    def putfo(self, fl, path, callback=None, confirm=True):
        request = self.parse_path(path)
        self.check_action(request, action=PERMS_ACTION_NAME_UPLOAD_FILE)

        client, rpath = self.get_sftp_client_rpath(request)
        success = False
        try:
            attr = client.putfo(fl, rpath, callback=callback, confirm=confirm)
            success = True
            return attr
        finally:
            self.create_ftp_log(path, 'Upload', success)

    def getfo(self, path, fl, callback=None):
        request = self.parse_path(path)
        self.check_action(request, action=PERMS_ACTION_NAME_DOWNLOAD_FILE)

        client, rpath = self.get_sftp_client_rpath(request)
        success = False
        try:
            attr = client.getfo(rpath, fl, callback=callback)
            success = True
            return attr
        finally:
            self.create_ftp_log(path, 'Download', success)

    def close(self):
        return self.session_ended()
